
import taichi as ti

ti.init(ti.cuda)

# global control  全局控制
paused = ti.field(ti.i32, ())

# number of planets 行星个数
N = 1000
# pos, vel and force of the planets   指定初始化的所有速度位置和受力  都为全局变量
# Nx2 vectors
pos = ti.Vector.field(2, ti.f32, N)
vel = ti.Vector.field(2, ti.f32, N)
force = ti.Vector.field(2, ti.f32, N)
mass = ti.field(ti.f32,N)
# radius = ti.Vector.field(1,ti.f32,N) 

@ti.data_oriented
class celestial_init:  
    # gravitational constant 6.67408e-11, using 1 for simplicity    重力常数   以及圆周率
    G = 1
    PI = 3.141592653
    # unit mass 行星质量
    m = 5
    # galaxy size  星系大小
    galaxy_size = 0.4
    # planet radius (for rendering)   行星可见大小
    planet_radius = 2
    # init vel  初始速度
    init_vel = 120
    # time-step size    事件步长
    h = 1e-5
    # substepping   微分分布时间t切块为10个dt
    substepping = 10

    #初始化 当前所有行星的速度位置和受力  散布于一个圆环
    @ti.kernel
    def initialize(self,num_st,num_ed):
        center=ti.Vector([0.5, 0.5])
        for i in range(num_st,num_ed):
            theta = ti.random() * 2 * self.PI
            r = (ti.sqrt(ti.random()) * 0.7 + 0.3) * self.galaxy_size
            offset = r * ti.Vector([ti.cos(theta), ti.sin(theta)])
            pos[i] = center+offset
            vel[i] = [-offset.y, offset.x]
            vel[i] *= self.init_vel 
            mass[i] = self.m
            # radius[i] = self.planet_radius


@ti.data_oriented
class celestial_compute: 

    # gravitational constant 6.67408e-11, using 1 for simplicity    重力常数   以及圆周率
    G = 1
    PI = 3.141592653
    # unit mass 行星质量
    m = 5
    # galaxy size   星系大小
    galaxy_size = 0.4
    # planet radius (for rendering)   行星可见大小
    planet_radius = 2
    # init vel  初始速度
    init_vel = 120
    # time-step size    事件步长
    h = 1e-5
    # substepping   微分分布时间t切块为10个dt
    substepping = 10
 
    #计算引力
    @ti.kernel
    def compute_force(self):

        # clear force  清空所有行星初速度
        for i in range(N):
            force[i] = ti.Vector([0.0, 0.0])

        # compute gravitational force
        for i in range(N):
            p = pos[i]
            for j in range(i): # bad memory footprint and load balance, but better CPU performance
                diff = p-pos[j]
                r = diff.norm(1e-5)

                # gravitational force -(GMm / r^2) * (diff/r) for i
                f = -self.G * mass[i] * mass[j] * (1.0/r)**3 * diff

                # assign to each particle
                force[i] += f
                force[j] += -f
            ''' 
            # for j in range(N):# double the computation for a better memory footprint and load balance
            #     if i != j: 
            #         diff = p-pos[j]
            #         r = diff.norm(1e-5)

            #         # gravitational force -(GMm / r^2) * (diff/r) for i
            #         f = -G * m * m * (1.0/r)**3 * diff

            #         # assign to each particle
            #         force[i] += f
            '''

    #根据每个单位时间受力的计算 更新速度量 和位置量
    @ti.kernel
    def update(self):
        dt = self.h/self.substepping
        for i in range(N):
            #symplectic euler
            vel[i] += dt*force[i]/self.m
            pos[i] += dt*vel[i]
 

@ti.data_oriented
class planet: 

    # gravitational constant 6.67408e-11, using 1 for simplicity    重力常数   以及圆周率
    G = 1
    PI = 3.141592653
    # unit mass 行星质量
    m = 5
    # galaxy size   星系大小
    galaxy_size = 0.4
    # planet radius (for rendering)   行星可见大小
    planet_radius = 2
    # init vel  初始速度
    init_vel = 120
    # time-step size    事件步长
    h = 1e-5
    # substepping   微分分布时间t切块为10个dt
    substepping = 10

    #初始化 当前所有行星的速度位置和受力  散布于一个圆环
    @ti.kernel
    def initialize(self):
        center=ti.Vector([0.5, 0.5])
        for i in range(N):
            theta = ti.random() * 2 * self.PI
            r = (ti.sqrt(ti.random()) * 0.7 + 0.3) * self.galaxy_size
            offset = r * ti.Vector([ti.cos(theta), ti.sin(theta)])
            pos[i] = center+offset
            vel[i] = [-offset.y, offset.x]
            vel[i] *= self.init_vel 

    #计算引力
    @ti.kernel
    def compute_force(self):

        # clear force  清空所有行星初速度
        for i in range(N):
            force[i] = ti.Vector([0.0, 0.0])

        # compute gravitational force
        for i in range(N):
            p = pos[i]
            for j in range(i): # bad memory footprint and load balance, but better CPU performance
                diff = p-pos[j]
                r = diff.norm(1e-5)

                # gravitational force -(GMm / r^2) * (diff/r) for i
                f = -self.G * self.m * self.m * (1.0/r)**3 * diff

                # assign to each particle
                force[i] += f
                force[j] += -f
            ''' 
            # for j in range(N):# double the computation for a better memory footprint and load balance
            #     if i != j: 
            #         diff = p-pos[j]
            #         r = diff.norm(1e-5)

            #         # gravitational force -(GMm / r^2) * (diff/r) for i
            #         f = -G * m * m * (1.0/r)**3 * diff

            #         # assign to each particle
            #         force[i] += f
            '''

    #根据每个单位时间受力的计算 更新速度量 和位置量
    @ti.kernel
    def update(self):
        dt = self.h/self.substepping
        for i in range(N):
            #symplectic euler
            vel[i] += dt*force[i]/self.m
            pos[i] += dt*vel[i]
 

#多体问题
gui = ti.GUI('3-body problem', (1080, 1080))

planet_body = celestial_init()
planet_body.m = 5
planet_body.planet_radius = 2
planet_body.initialize(0,N-3)

solar_body = celestial_init()
solar_body.m= 50
solar_body.planet_radius=6
solar_body.initialize(N-3,N)

computer = celestial_compute()

while gui.running:

    for i in range(computer.substepping):
        computer.compute_force()
        computer.update()

    gui.clear(0x112F41)
    gui.circles(pos.to_numpy()[0:N-3], color=0xffff6f, radius=2)
    gui.circles(pos.to_numpy()[N-3:N], color=0x0f8f6f, radius=8)
    gui.show()